"""
Name/Country 数据集,本例中对变长的name size进行补齐,使用最长的name size作为标准size
"""
import csv
import gzip
import heapq

from torch.utils.data import Dataset


class Name2CountryDatasetWithComplement(Dataset):
    def __init__(self, isTrain: bool):
        """
        读取数据,根据输入的isTrain判定加载的数据集类型
        :param isTrain: True表示读取train的数据,False为test数据
        """
        super().__init__()
        # 根据输入判定当前的数据集类型
        filePath = "./dataset/names_train.csv.gz" if isTrain else "./dataset/names_test.csv.gz"
        # 加载数据,使用rt打开(以文本类型进行读取)
        with gzip.open(filePath, "rt") as f:
            # 使用csv读取文件
            reader = csv.reader(f)
            # 读取所有的信息
            rawData = list(reader)
            # 将信息分离成name/country
            self.names, self.countries = zip(*rawData)
            # 记录样本总数
            self.len = len(self.names)
            # 建立country索引表
            self.countries_list = list(sorted(set(self.countries)))
            # 建立country:idx词典
            self.countries_dict = {v: idx for idx, v in enumerate(self.countries_list)}
            # country数
            self.countries_num = len(self.countries_list)
            # 最大的name
            # bigName = heapq.nlargest(1, self.names, key=lambda x: len(x))
            # 获得最大name的长度
            # nameLength = len(bigName[0])
            # 将此长度作为标准长度建立表

    def __getitem__(self, index):
        # 返回name,country idx
        return self.names[index], self.countries_dict[self.countries[index]]

    def __len__(self):
        return self.len

    def idx2Country(self, idx):
        # 通过索引返回country字符串
        return self.countries_dict[idx]


if __name__ == '__main__':
    dataset = Name2CountryDatasetWithComplement(False)
